

package net.qiqbframework.messaging.task;


import lombok.extern.slf4j.Slf4j;
import net.qiqbframework.common.Assert;
import net.qiqbframework.messaging.Message;
import net.qiqbframework.messaging.MetaData;
import net.qiqbframework.messaging.correlation.CorrelationDataProvider;

import java.util.*;
import java.util.function.Consumer;


@Slf4j
public abstract class AbstractMessageTask<T extends Message<?>> implements MessageTask<T> {
    /**
     * 任务处理过程中， 需要使用的一些资源。这些资源在任务执行后清除
     */
    private final Map<String, Object> resources = new HashMap<>();

    private final Collection<CorrelationDataProvider> correlationDataProviders = new LinkedHashSet<>();
    private MessageTask<?> parentMessageWork;
    private TaskPhase phase = TaskPhase.NOT_STARTED;
    private boolean rolledBack;

    @Override
    public void start() {
        if (log.isDebugEnabled()) {
            log.debug("Message Task Starting ...");
        }
        Assert.state(TaskPhase.NOT_STARTED.equals(phase()), () -> "Message Task is already started");
        rolledBack = false;
        onRollback(u -> rolledBack = true);
        CurrentMessageTask.ifStarted(parent -> {
            // we're nesting.
            this.parentMessageWork = parent;
            root().onCleanup(r -> changePhase(TaskPhase.CLEANUP, TaskPhase.CLOSED));
        });
        changePhase(TaskPhase.STARTED);
        CurrentMessageTask.set(this);
    }

    @Override
    public void plan() {
        changePhase(TaskPhase.PLAN);
    }

    @Override
    public void onPlan(Consumer<MessageTask<T>> handler) {
        addHandler(TaskPhase.PLAN, handler);
    }

    @Override
    public void commit() {

        //Assert.state(phase() == TaskPhase.PLAN, () -> String.format("The MessageWork is in an incompatible phase: %s", phase()));
        Assert.state(isCurrent(), () -> "The MessageWork is not the current Unit of Work");
        try {
            if (isRoot()) {
                commitAsRoot();
            } else {
                commitAsNested();
            }
        } finally {
            CurrentMessageTask.clear(this);
        }
    }

    private void commitAsRoot() {
        try {
            try {
                changePhase(TaskPhase.PREPARE_COMMIT, TaskPhase.COMMIT);
            } catch (Exception e) {
                setRollbackCause(e);
                changePhase(TaskPhase.ROLLBACK);
                throw e;
            }
            if (phase() == TaskPhase.COMMIT) {
                changePhase(TaskPhase.AFTER_COMMIT);
            }
        } finally {
            changePhase(TaskPhase.CLEANUP, TaskPhase.CLOSED);
        }
    }

    private void commitAsNested() {
        try {
            changePhase(TaskPhase.PREPARE_COMMIT, TaskPhase.COMMIT);
            delegateAfterCommitToParent(this);
            parentMessageWork.onRollback(u -> changePhase(TaskPhase.ROLLBACK));
        } catch (Exception e) {
            setRollbackCause(e);
            changePhase(TaskPhase.ROLLBACK);
            throw e;
        }
    }

    private void delegateAfterCommitToParent(MessageTask<?> uow) {
        Optional<MessageTask<?>> parent = uow.parent();
        if (parent.isPresent()) {
            parent.get().afterCommit(this::delegateAfterCommitToParent);
        } else {
            changePhase(TaskPhase.AFTER_COMMIT);
        }
    }

    @Override
    public void rollback(Throwable cause) {
        if (log.isDebugEnabled()) {
            log.debug("Rolling back Unit Of Work.", cause);
        }
        Assert.state(isActive() && phase().isBefore(TaskPhase.ROLLBACK),
                () -> String.format("The MessageWork is in an incompatible phase: %s", phase()));
        Assert.state(isCurrent(), () -> "The MessageWork is not the current Unit of Work");
        try {
            setRollbackCause(cause);
            changePhase(TaskPhase.ROLLBACK);
            if (isRoot()) {
                changePhase(TaskPhase.CLEANUP, TaskPhase.CLOSED);
            }
        } finally {
            CurrentMessageTask.clear(this);
        }
    }

    @Override
    public Optional<MessageTask<?>> parent() {
        return Optional.ofNullable(parentMessageWork);
    }

    @Override
    public Map<String, Object> resources() {
        return resources;
    }

    @Override
    public boolean isRolledBack() {
        return rolledBack;
    }

    @Override
    public void registerCorrelationDataProvider(CorrelationDataProvider correlationDataProvider) {
        correlationDataProviders.add(correlationDataProvider);
    }

    @Override
    public MetaData getCorrelationData() {
        if (correlationDataProviders.isEmpty()) {
            return MetaData.emptyInstance();
        }
        Map<String, Object> result = new HashMap<>();
        for (CorrelationDataProvider correlationDataProvider : correlationDataProviders) {
            try {
                final Map<String, ?> extraData = correlationDataProvider.correlationDataFor(getMessage());
                if (extraData != null) {
                    result.putAll(extraData);
                }
            } catch (Exception e) {
                log.warn(
                        "Encountered exception creating correlation data for message with id: '{}' "
                                + "using correlation provider with class: '{}'"
                                + "will continue without as this might otherwise prevent a rollback.",
                        getMessage().getIdentifier(),
                        correlationDataProvider.getClass(),
                        e);
            }
        }
        return MetaData.of(result);
    }

    @Override
    public void onPrepareCommit(Consumer<MessageTask<T>> handler) {
        addHandler(TaskPhase.PREPARE_COMMIT, handler);
    }

    @Override
    public void onCommit(Consumer<MessageTask<T>> handler) {
        addHandler(TaskPhase.COMMIT, handler);
    }

    @Override
    public void afterCommit(Consumer<MessageTask<T>> handler) {
        addHandler(TaskPhase.AFTER_COMMIT, handler);
    }

    @Override
    public void onRollback(Consumer<MessageTask<T>> handler) {
        addHandler(TaskPhase.ROLLBACK, handler);
    }

    @Override
    public void onCleanup(Consumer<MessageTask<T>> handler) {
        addHandler(TaskPhase.CLEANUP, handler);
    }

    @Override
    public TaskPhase phase() {
        return phase;
    }

    /**
     * Overwrite the current phase with the given {@code phase}.
     *
     * @param phase the new phase of the Unit of Work
     */
    protected void setPhase(TaskPhase phase) {
        this.phase = phase;
    }


    protected void changePhase(TaskPhase... phases) {
        for (TaskPhase phase : phases) {
            setPhase(phase);
            notifyHandlers(phase);
        }
    }


    protected Collection<CorrelationDataProvider> correlationDataProviders() {
        return correlationDataProviders;
    }

    protected abstract void notifyHandlers(TaskPhase phase);

    protected abstract void addHandler(TaskPhase phase, Consumer<MessageTask<T>> handler);

    /**
     * Set the execution result of processing the current {@link #getMessage() Message}.
     *
     * @param executionResult the ExecutionResult of the currently handled Message
     */
    protected abstract void setExecutionResult(ExecutionResult executionResult);

    /**
     * Sets the cause for rolling back this Unit of Work.
     *
     * @param cause The cause for rolling back this Unit of Work
     */
    protected abstract void setRollbackCause(Throwable cause);
}
